package llmserviceopenai.tools

import llmserviceopenai.model.ChatCompletionRequest
import llmserviceopenai.model.EmbeddingRequest
import llmserviceopenai.model.EmbeddingModelResponse
import llmserviceopenai.model.ChatMessage

import ragserverengine.milvus.MilvusRequest
import ragserverengine.milvus.MilvusSearchRequest
import ragserverengine.milvus.MilvusSearchResponse

import kotlinx.serialization.json.Json
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.*

interface Tool {

    val name: String
    suspend fun execute(input: String): String
}

class DbTool : Tool {

    private val tool_dist = "This is Vertor Data Base Tool"

    override val name: String
        get() = tool_dist

    override suspend fun execute(input: String): String {
        return "DbTool return is xxxxx"
    }
}

fun ChatCompletionRequest.toPayloadJson(): JsonObject {
    // 先序列化整个对象为 JsonObject
    val jsonMap = Json.encodeToJsonElement(this).jsonObject.toMutableMap()

    // 展开 extra_body 到顶层
    jsonMap.remove("extra_body")?.jsonObject?.forEach { (key, value) ->
        if (!jsonMap.containsKey(key)) {
            jsonMap[key] = value
        }
    }

    return JsonObject(jsonMap)
}

// chat请求转换成embedding请求格式
fun ChatCompletionRequest.toEmbeddingRequest(embeddingModel: String): EmbeddingRequest {
    // 提取最后一个 user 消息的内容
    val lastUserQuery = messages.asReversed().find { it.role == "user" }?.content
        ?: throw IllegalArgumentException("No content message found in chat request")

    return EmbeddingRequest(
        model = embeddingModel, // 可以指定专用 embedding 模型，如"text-embedding-ada-002"
        input = listOf(lastUserQuery) // List<String>
    )
}

fun EmbeddingModelResponse.toMilvusSearchRequest(
    collectionName: String,
    topK: Int = 10,
    metricType: String = "COSINE", // 推荐使用 COSINE，如果你的 Milvus 索引也是基于 cosine
    outputFields: List<String>? = listOf("id", "doc", "meta")
): MilvusSearchRequest {
    // 提取所有 embedding 向量，包装成 List<List<Float>>
    val vectors = data.map { it.embedding }

    return MilvusSearchRequest(
        collectionName = collectionName,
        vector = vectors,           // 支持 batch search（多个 query 向量）
        topK = topK,
        metricType = metricType,
        outputFields = outputFields
    )
}

private fun JsonElement.toKotlinValue(): Any? = when (this) {
    is JsonNull -> null
    is JsonPrimitive -> when {
        isString -> content
        booleanOrNull != null -> boolean
        longOrNull != null -> long
        doubleOrNull != null -> double
        else -> content
    }

    is JsonArray -> map { it.toKotlinValue() }
    is JsonObject -> mapValues { it.value.toKotlinValue() }

}


/**
 * 构建带上下文（RAG）的 ChatCompletionRequest
 *
 * @param originalRequest 用户原始请求
 * @param milvusResponse Milvus 检索结果
 * @return 增强后的 ChatCompletionRequest，可直接用于大模型调用
 */
fun buildRagRequest(
    originalRequest: ChatCompletionRequest,
    milvusResponse: MilvusSearchResponse
): ChatCompletionRequest {

    // 提取 Milvus 检索出的内容
    val contexts = milvusResponse.results.firstOrNull()
        ?.map { it.doc }
        ?.filterNotNull()
        ?.filterNot { it.isBlank() }
        ?: emptyList()

    // 拼接上下文
    val contextText = if (contexts.isEmpty()) {
        ""
    } else {
        contexts.withIndex().joinToString("\n") { (index, content) ->
            "[${index + 1}] $content"
        }
    }

    // 提取用户原始问题
    val originalQuery = originalRequest.messages.lastOrNull { it.role == "user" }?.content
        ?: "请回答以下问题"

    // 构造新的消息列表
    val augmentedMessages = buildList {
        // 系统提示词：嵌入上下文
        add(
            ChatMessage(
                role = "system",
                content = """
                    你是一个智能助手。以下是从知识库检索到的参考信息，请在回答时尽量参考这些内容（如果相关）：
                    -----
                    $contextText
                    -----
                    如果参考内容不相关，你可以忽略它们并直接回答用户问题。
                """.trimIndent()
            )
        )

        // 可以保留原始上下文（可选）
        addAll(originalRequest.messages)

        // 确保用户的 query 在最后
        add(ChatMessage(role = "user", content = originalQuery))
    }

    // 返回增强后的请求
    return originalRequest.copy(
        messages = augmentedMessages
    )
}